package test.ground.kafka;

import com.wang.tool.monitor.core.MonitorInfo;
import com.wang.tool.monitor.template.AbsMonitorInvoker;
import com.wang.tool.trace.TraceContext;
import com.wang.tool.trace.TraceHolder;
import com.wang.tool.trace.TraceTemplate;
import org.apache.kafka.clients.producer.ProducerInterceptor;
import org.apache.kafka.clients.producer.ProducerRecord;
import org.apache.kafka.clients.producer.RecordMetadata;

import java.util.Map;
import java.util.function.BiConsumer;

public class KafkaProducerInterceptor implements ProducerInterceptor<String, String> {  

    @Override  
    public ProducerRecord<String, String> onSend(ProducerRecord<String, String> record) {  
        // Modify the record if needed
        return TraceTemplate.monitorWithNext("mq", "kafka", record.topic(), new AbsMonitorInvoker<ProducerRecord<String, String>>() {
            @Override
            public ProducerRecord<String, String> internalCall() throws Throwable {
                return record;
            }
        }, new BiConsumer<TraceContext, MonitorInfo<ProducerRecord<String, String>>>() {
            @Override
            public void accept(TraceContext context, MonitorInfo<ProducerRecord<String, String>> producerRecordMonitorInfo) {
                if (context!=null) {
                    record.headers().add(TraceHolder.TRACE_ID_HEADER, context.getTraceId().getBytes());
                    record.headers().add(TraceHolder.RPC_ID_HEADER, context.getRpcId().getBytes());
                }
            }
        });
    }  

    @Override  
    public void onAcknowledgement(RecordMetadata metadata, Exception exception) { }  

    @Override  
    public void close() { }  

    @Override  
    public void configure(Map<String, ?> configs) { }  
}